import os
import xml.etree.ElementTree as ET
import json
from typing import Dict, Any, Generator

LOCAL_JSON_PATH = r"C:\local_pubmed_json"

def load_large_xml_data_to_json(file_path: str) -> Generator[Dict, None, None]:
    """读取XML文件 Args:file_path: XML文件路径  Yields: Dict: 解析后的数据项  """
    try:
        for event, elem in ET.iterparse(file_path, events=('start', 'end')):
            if event == 'end' and elem.tag == 'PubmedArticle':
                parsed_article = parse_pubmed_article(elem)
                pmid = parsed_article.get('PMID')
                year = parsed_article.get('Year')
                article_title = parsed_article.get('ArticleTitle')
                title = parsed_article.get('JournalTitle')
                abstract_text = parsed_article.get('AbstractText')
                try:
                    if (pmid and int(pmid) > 31000000 and title
                             and article_title and year and abstract_text):
                        title_display = article_title[:50] if article_title else 'N/A'
                        print(f"解析文章: PMID={pmid},{ year}, {title} ,标题={title_display}...")
                        save_article_to_json(parsed_article, LOCAL_JSON_PATH)
                        yield parsed_article
                except (ValueError, TypeError):
                    pass
                elem.clear()
    except ET.ParseError as e:
        print(f"XML解析错误: {e}")
    except FileNotFoundError:
        print(f"文件未找到: {file_path}")
    except MemoryError:
        print("内存不足，无法处理大文件")

def save_article_to_json(article_data: Dict[str, Any], base_path: str):
    """将文章数据保存为JSON文件，按PMID路径分割"""
    pmid = article_data.get('PMID')
    if not pmid:
        print("PMID不存在，跳过保存")
        return
    path_prefix_4 = pmid[:4] if len(pmid) >= 4 else pmid
    path_suffix_3 = pmid[-3:] if len(pmid) >= 3 else pmid
    dir_path = os.path.join(base_path, path_prefix_4, path_suffix_3)
    os.makedirs(dir_path, exist_ok=True)
    file_path = os.path.join(dir_path, f"{pmid}.json")
    # 如果文件已存在，则跳过保存
    if os.path.exists(file_path):
        print(f"文章 {pmid} 已存在，跳过保存")
        return
    with open(file_path, 'w', encoding='utf-8') as f:
        json.dump(article_data, f, ensure_ascii=False, indent=2)
    print(f"文章 {pmid} 已保存至 {file_path}")


def parse_pubmed_article(element) -> Dict[str, Any]:
    """
    解析单个PubMed文章元素

    Args:
        element: XML元素

    Returns:
        Dict: 解析后的文章数据
    """
    data = {}

    # 解析MedlineCitation部分
    medline_citation = element.find('MedlineCitation')
    if medline_citation is not None:
        # 提取PMID
        pmid_elem = medline_citation.find('PMID')
        if pmid_elem is not None:
            data['PMID'] = pmid_elem.text
            data['PMID_Version'] = pmid_elem.get('Version')

        # 提取DateCompleted
        date_completed = medline_citation.find('DateCompleted')
        if date_completed is not None:
            year_elem = date_completed.find('Year')
            month_elem = date_completed.find('Month')
            day_elem = date_completed.find('Day')
            parts = []
            if year_elem is not None and year_elem.text:
                parts.append(year_elem.text)
            if month_elem is not None and month_elem.text:
                parts.append(month_elem.text)
            if day_elem is not None and day_elem.text:
                parts.append(day_elem.text)
            data['DateCompleted'] = "-".join(parts)

        # 提取DateRevised
        date_revised = medline_citation.find('DateRevised')
        if date_revised is not None:
            year_elem = date_revised.find('Year')
            month_elem = date_revised.find('Month')
            day_elem = date_revised.find('Day')
            parts = []
            if year_elem is not None and year_elem.text:
                parts.append(year_elem.text)
            if month_elem is not None and month_elem.text:
                parts.append(month_elem.text)
            if day_elem is not None and day_elem.text:
                parts.append(day_elem.text)
            data['DateRevised'] = "-".join(parts)

        # 提取文章信息
        article = medline_citation.find('Article')
        if article is not None:
            # 提取期刊信息
            journal = article.find('Journal')
            if journal is not None:
                # ISSN信息
                issn = journal.find('ISSN')
                if issn is not None:
                    data['ISSN'] = issn.text
                    data['ISSN_Type'] = issn.get('IssnType')

                # JournalIssue信息
                journal_issue = journal.find('JournalIssue')
                if journal_issue is not None:
                    # Volume
                    volume = journal_issue.find('Volume')
                    if volume is not None:
                        data['Volume'] = volume.text

                    # Issue
                    issue = journal_issue.find('Issue')
                    if issue is not None:
                        data['Issue'] = issue.text

                    # PubDate
                    pub_date = journal_issue.find('PubDate')
                    if pub_date is not None:
                        pub_year = pub_date.find('Year')
                        pub_month = pub_date.find('Month')
                        if pub_year is not None:
                            data['PubDate_Year'] = pub_year.text
                        if pub_month is not None:
                            data['PubDate_Month'] = pub_month.text

                # 期刊标题
                journal_title = journal.find('Title')
                if journal_title is not None:
                    data['JournalTitle'] = journal_title.text

                # ISOAbbreviation
                iso_abbr = journal.find('ISOAbbreviation')
                if iso_abbr is not None:
                    data['ISOAbbreviation'] = iso_abbr.text

            # 文章标题
            title_elem = article.find('ArticleTitle')
            if title_elem is not None:
                data['ArticleTitle'] = title_elem.text
            # 提取摘要
            abstract = article.find('Abstract')
            if abstract is not None:
                abstract_text = abstract.find('AbstractText')
                if abstract_text is not None:
                    data['AbstractText'] = abstract_text.text

            # 提取作者信息
            author_list = article.find('AuthorList')
            if author_list is not None:
                authors = []
                for author in author_list.findall('Author'):
                    last_name = author.find('LastName')
                    fore_name = author.find('ForeName')
                    initials = author.find('Initials')
                    author_info = {}
                    if last_name is not None:author_info['LastName'] = last_name.text
                    if fore_name is not None:author_info['ForeName'] = fore_name.text
                    if initials is not None:author_info['Initials'] = initials.text

                    if author_info:
                        authors.append(author_info)
                data['Authors'] = authors

            # 提取语言
            language = article.find('Language')
            if language is not None:data['Language'] = language.text

            # 提取Pagination
            pagination = article.find('Pagination')
            if pagination is not None:
                medline_pgn = pagination.find('MedlinePgn')
                if medline_pgn is not None:data['MedlinePgn'] = medline_pgn.text

    # 解析PubmedData部分
    pubmed_data = element.find('PubmedData')
    if pubmed_data is not None:
        # 提取History
        history = pubmed_data.find('History')
        if history is not None:
            pub_dates = []
            for pub_date in history.findall('PubMedPubDate'):
                pub_date_dict = {
                    'PubStatus': pub_date.get('PubStatus'),
                }
                year = pub_date.find('Year')
                month = pub_date.find('Month')
                day = pub_date.find('Day')
                hour = pub_date.find('Hour')
                minute = pub_date.find('Minute')

                if year is not None:pub_date_dict['Year'] = year.text
                if month is not None:pub_date_dict['Month'] = month.text
                if day is not None:pub_date_dict['Day'] = day.text
                if hour is not None: pub_date_dict['Hour'] = hour.text
                if minute is not None:pub_date_dict['Minute'] = minute.text

                pub_dates.append(pub_date_dict)
            data['PubMedPubDateHistory'] = pub_dates
        # 提取ArticleIdList
        article_id_list = pubmed_data.find('ArticleIdList')
        if article_id_list is not None:
            article_ids = []
            for article_id in article_id_list.findall('ArticleId'):
                article_ids.append({
                    'IdType': article_id.get('IdType'),
                    'Value': article_id.text
                })
            data['ArticleIds'] = article_ids
        # 提取PublicationStatus
        publication_status = pubmed_data.find('PublicationStatus')
        if publication_status is not None:
            data['PublicationStatus'] = publication_status.text
    return data

